import torch

def print_torch_ckpt_params(ckpt_path: str):
    print(f"===== PyTorch ckpt 参数 ({ckpt_path}) =====")
    state_dict = torch.load(ckpt_path, map_location="cpu")

    # 如果是 nn.Module
    if isinstance(state_dict, torch.nn.Module):
        state_dict = state_dict.state_dict()

    # 如果是常见字典格式 { "epoch":..., "model":..., "optimizer":... }
    if isinstance(state_dict, dict) and "model" in state_dict:
        print(f"检测到 ckpt 包含训练信息 (epoch, optimizer)，进入 model 部分")
        state_dict = state_dict["model"]

    # 打印参数
    for name, param in state_dict.items():
        if isinstance(param, torch.Tensor):
            print(f"{name}: {tuple(param.shape)}")
        else:
            print(f"{name}: (非Tensor参数，类型={type(param)})")


if __name__ == "__main__":
    print_torch_ckpt_params("model_000014.ckpt")
